module Random = struct
(* generate a random permutation of {1,...,n} 
   following the Knuth shuffle algorithm *)
let randperm ?rnd_state n =
  if ( n<= 0 ) then
    raise (Invalid_argument "randperm : n <= 0" )
  else
    let v = Array.init n (fun i -> i+1) in
    let state = 
      match rnd_state with
      | None -> Random.get_state ()
      | Some state -> state in
    for i = 0 to (n-2) do
      let j = (Random.State.int state (n-i)) + i in
      let a = v.(j) in
        v.(j) <- v.(i);
        v.(i) <- a;
    done;
    if rnd_state = None then Random.set_state state;
    v;;

(* generate a random number according to the normal law
   with mean mu and variance sigma^2 
   following M. Wichura AS241 algorithm *)
let randn ?rnd_state mu sigma =
  let state = 
    match rnd_state with
    | None -> Random.get_state ()
    | Some state -> state in
  let probit x =
    let zero = 0.0 and one = 1.0 and half = 0.5
    and split1 = 0.425 and split2 = 5.0
    and const1 = 0.180625 and const2 = 1.6 in
    let a = 
    [|3.3871328727963666080e0;
    1.3314166789178437745e+2;
    1.9715909503065514427e+3;
    1.3731693765509461125e+4;
    4.5921953931549871457e+4;
    6.7265770927008700853e+4;
    3.3430575583588128105e+4;
    2.5090809287301226727e+3|]
    and b = 
    [|0.0;
    4.2313330701600911252e+1;
    6.8718700749205790830e+2;
    5.3941960214247511077e+3;
    2.1213794301586595867e+4;
    3.9307895800092710610e+4;
    2.8729085735721942674e+4;
    5.2264952788528545610e+3|]
    and c =
    [|1.42343711074968357734e0;
    4.63033784615654529590e0;
    5.76949722146069140550e0;
    3.64784832476320460504e0;
    1.27045825245236838258e0;
    2.41780725177450611770e-1;
    2.27238449892691845833e-2;
    7.74545014278341407640e-4|]
    and d =
    [|0.0;
    2.05319162663775882187e0;
    1.67638483018380384940e0;
    6.89767334985100004550e-1;
    1.48103976427480074590e-1;
    1.51986665636164571966e-2;
    5.47593808499534494600e-4;
    1.05075007164441684324e-9|]
    and e =
    [|6.65790464350110377720e0;
    5.46378491116411436990e0;
    1.78482653991729133580e0;
    2.96560571828504891230e-1;
    2.65321895265761230930e-2;
    1.24266094738807843860e-3;
    2.71155556874348757815e-5;
    2.01033439929228813265e-7|]
    and f =
    [|0.0;
    5.99832206555887937690e-1;
    1.36929880922735805310e-1;
    1.48753612908506148525e-2;
    7.86869131145613259100e-4;
    1.84631831751005468180e-5;
    1.42151175831644588870e-7;
    2.04426310338993978564e-15|] in
    let q = x -. half in
    if ( abs_float q <= split1 ) then
      let r = const1 -. q *. q in
      q *. (((((((a.(7) *. r +. a.(6)) *. r +. a.(5)) *. r +. a.(4)) *. r +. a.(3))
      *. r +. a.(2)) *. r +. a.(1)) *. r +. a.(0)) /.
      (((((((b.(7) *. r +. b.(6)) *. r +. b.(5)) *. r +. b.(4)) *. r +. b.(3))
      *. r +. b.(2)) *. r +. b.(1)) *. r +. one)
    else
      let r = if ( q < zero ) then x else one -. x in
      if ( r <= 0.0 ) then 0.0
      else
        let ret =
          let r' = sqrt( -.log(r) ) in
          if ( r' <= split2 ) then
            let r'' = r' -. const2 in
            (((((((c.(7) *. r'' +. c.(6)) *. r'' +. c.(5)) *. r'' +. c.(4)) *. r'' +. c.(3))
            *. r'' +. c.(2)) *. r'' +. c.(1)) *. r'' +. c.(0)) /.
            (((((((d.(7) *. r'' +. d.(6)) *. r'' +. d.(5)) *. r'' +. d.(4)) *. r'' +. d.(3))
            *. r'' +. d.(2)) *. r'' +. d.(1)) *. r'' +. one)
          else
            let r'' = r' -. split2 in
            (((((((e.(7) *. r'' +. e.(6)) *. r'' +. e.(5)) *. r'' +. e.(4)) *. r'' +. e.(3))
            *. r'' +. e.(2)) *. r'' +. e.(1)) *. r'' +. e.(0)) /.
            (((((((f.(7) *. r'' +. f.(6)) *. r'' +. f.(5)) *. r'' +. f.(4)) *. r'' +. f.(3))
            *. r'' +. f.(2)) *. r'' +. f.(1)) *. r'' +. one) in
        if ( q < 0.0 ) then
          -.ret
        else
          ret
  in
    let x = mu +. sigma *. (probit (Random.State.float state 1.0)) in
    begin
      if rnd_state = None then Random.set_state state;
      x
    end;;

end;;

module Float = struct

let sign (x:float) =
   if ( x > 0.0 ) then
      1.0
   else if ( x < 0.0 ) then
      -1.0
   else
     0.0;;

let max (x:float) (y:float) =
   if ( x >= y ) then
      x
   else
      y;;

let min (x:float) (y:float) =
   if ( x <= y ) then
      x
   else
      y;;

end;;

module NonLinear = struct

(* S_{\lambda}(c) := \min_x \frac{1}{2}( c - x )^2 + \lambda | x |
                   = sgn( c ) * max( 0, |c| - \lambda ) *)
let soft_threshold lambda x =
   (Float.sign x) *. (Float.max 0.0 ((abs_float x) -. lambda));;

end;;
